# encoding=utf-8
import os
from seg_common.BatchLoader import BatchLoaderFactory
from seg_common.annotation import CommonAnnotation
from torchvision import transforms
from PIL import Image
from seg_system import ApplicationConfig


class CSNetBatchFactory(BatchLoaderFactory):
    def __init__(self):
        super(CSNetBatchFactory, self).__init__()
        self.from_path = ""
        self.idx_to_token, self.token_to_idx = [], dict()
        self.transform = transforms.Compose([
            transforms.Resize([384, 384]),
            transforms.ToTensor()
        ])

    @CommonAnnotation.override()
    def load_data(self, path: str, **kwargs):
        self.idx_to_token.clear()
        self.token_to_idx.clear()
        self.from_path = path
        for each_file in os.listdir(path):
            self.idx_to_token.append(each_file)
            self.token_to_idx[each_file] = len(self.idx_to_token) - 1

    @CommonAnnotation.override()
    def prepare_data(self, index: int):
        o = self.idx_to_token[index]
        obj = os.path.join(self.from_path, o)
        image = Image.open(obj)
        gray = image.convert('L')
        trans = self.transform(gray)
        trans = trans.to(ApplicationConfig.SystemConfig.DEVICE)
        index = self.get_index(o)
        return trans, index

    @CommonAnnotation.override()
    def get_data_lens(self):
        return len(self.idx_to_token)

    @CommonAnnotation.override()
    def get_index(self, tokens):
        if not isinstance(tokens, (list, tuple)):
            return self.token_to_idx.get(tokens, 0)
        return [self.get_index(token) for token in tokens]

    @CommonAnnotation.override()
    def get_tokens(self, indices):
        if not isinstance(indices, (list, tuple)):
            return self.idx_to_token[indices]
        return [self.idx_to_token[index] for index in indices]
